#!/usr/bin/env python

import csv
from Bio import AlignIO
from Bio import SeqIO
from Bio.PDB import *
from Bio.SeqRecord import SeqRecord
from Bio.Seq import MutableSeq as Seq
from Bio.Align import MultipleSeqAlignment
import os
import re
import sys
#import MySQLdb
import sqlite3
import hmmer3
import pickle
import pprint
#from pprint import pprint
#INPUT: PDB files
#
#1. renumber all residues
#2. load a contacting chain file
#3. Calculate contacting residues from

def load_cc(filename):
    #INPUT: File name (str)
    #OUTPUT: List of lists
    holder = []

    with open(filename, 'rb') as f:
        reader = csv.reader(f, skipinitialspace=True)
        try:
            for row in reader:
                holder.append(row)
        except csv.Error, e:
            sys.exit('file %s, line %d: %s' % (filename, reader.line_num, e))
    return holder

def renum(structure, chain, from_to = False):
    #renumber residues in structure[0][chain]
    
    
    if from_to:
        (start, end) = from_to
        #renumber 
        i = 1
        for residue in structure[0][chain]:
            rid = residue.id
            if (rid[0] == ' ' and
                rid[1] in range(start,end+1)):
                
                residue.id = (' ', i, ' ')
                i +=1
            else:
                structure[0][chain].detach_child(rid)
    else:
        i = 1
        #print chain
        for residue in structure[0][chain]:
            rid = residue.id
            
            if (residue.id[0] == ' '):
                #print residue
                residue.id = (' ', i, ' ')
                #print residue
                i += 1
            else:
                structure[0][chain].detach_child(rid)
    return structure

def contact(structure, scid, tcid):
    #INPUT: PDB object, source chain, target chain
    #OUTPUT: dictionary: {(target_resid, source_resid): min distance}
    #Calculates the minimum distance of two residues in structure
    
    atom_list = Selection.unfold_entities(structure[0][scid], 'A')
    ns = NeighborSearch(atom_list)
    target_chain = structure[0][tcid]
    dist = dict()
    for res in target_chain:
        
        if res.id[0] == ' ':
            #ignore hetero residues
            iscontact = False
            for atom in res:
                if atom.name not in ['N', 'O', 'C']:
                    center = atom.get_coord()
                    neighbors = ns.search(center, 15.0)
                    if len(neighbors) != 0:
                        for r_atom in neighbors:
                            if r_atom.name not in ['N', 'O', 'C']:
                                rres = r_atom.get_parent()
                                
                                if rres.id[0] == ' ':
                                    #ignore hetero residue
                                    k = (res.id[1], rres.id[1])
                                    distance = atom - r_atom
                                    if dist.has_key(k):
                                        dist[k] = min(dist[k], distance)
                                    else:
                                        dist[k] = distance
    lastres = sorted(dist.keys())[-1][0]
    
    newdist = {}
    for key in dist:
        if key[0] > lastres - 5:
            newdist[(lastres - key[0], key[1])] = dist[key]
    #pprint.pprint(newdist)
    return newdist
                             

def hmmer_call(acc, hmmdb, seq, eval):
    #run hmmsearch on sequence seq
    
    tmpPath = "scanHmmer.tmp"
    tmpFile = open( tmpPath, "w" )
    #print "seq:", seq
    #ss = Seq(seq)
    record = SeqRecord(seq, acc, "", "")
    SeqIO.write(record, tmpFile, "fasta")
    tmpFile.close()
    
    try:
        cmdArray = ('hmmsearch -E'+str(eval)+' '+hmmdb+' '+tmpPath+
                    '> tmp.hmmscan')
        #print cmdArray
        os.system(str(cmdArray))
    except OSError:
        os.unlink( tmpPath )
        raise OSError( "Unable to find hmmsearch" )
    hmmscan_out = open('tmp.hmmscan', 'r')
    _return = hmmer3.parse(hmmscan_out)

    try:
        os.unlink( tmpPath )
    except OSError:
        pass
    #print _return
    return _return


def get_offset(targetseq, hmmseq, srcindex, alifrom, hmmfrom, glgf, gd):
    #convert the residue index from sequence index (res# in a sequence)
    #to motif + offset format
    #
    def transindex(_seq, _index, _from):
        #find (motif + offset)
        _pointer = 0
        _rescounter = 0
        _index = _index - _from
        while (_rescounter < _index or
               _seq[_pointer] == '.' or
               _seq[_pointer] =='-'):
            
            if _seq[_pointer] != '.' and _seq[_pointer] != '-':
                _rescounter += 1
            _pointer += 1
        return _pointer
    
    src_dex = transindex (targetseq, srcindex, alifrom)
    glgf_dex = transindex (hmmseq, glgf, hmmfrom)
    gd_dex = transindex (hmmseq, gd, hmmfrom)
    #print targetseq
    if hmmseq[src_dex] =='.':
        return False, False
    residue = targetseq[src_dex]
    gaps = 0 
    if src_dex > glgf_dex and src_dex < gd_dex:
        for r in hmmseq[glgf_dex:src_dex]:
            if r == '.' or r == '-':
                gaps += 1
        if src_dex - glgf_dex - gaps <= 14:
            return ('GLGF', src_dex - glgf_dex - gaps), residue
        else:
            gaps = 0
            for r in hmmseq[src_dex:gd_dex]:
                if r == '.' or r == '-':
                    gaps += 1
            #print src_dex, gd_dex, gaps
            return ('GD', src_dex - gd_dex + gaps), residue
    elif src_dex < glgf_dex:
        gaps = 0
        for r in hmmseq[src_dex:glgf_dex]:
            if r == '.' or r == '-':
                    gaps += 1
        return ('GLGF', src_dex - glgf_dex + gaps), residue
    elif src_dex > gd_dex:
        gaps = 0
        for r in hmmseq[gd_dex:src_dex]:
                if r == '.' or r == '-':
                    gaps += 1
        #print src_dex, gd_dex, gaps, residue
        return ('GD', src_dex - gd_dex - gaps), residue
    elif src_dex == glgf_dex:
        return ('GLGF', 0), residue
    elif src_dex == gd_dex:
        return ('GD', 0), residue
    
'''
final_out = {
    0: {(GLGF,1): [[],[],[]]
        (GLGF,2): [[],[],[]]},
    1: {},
    2: {},
    3: {},
    4: {},
}
'''

def find_motif(outdict, hmmout, dict_contact, peppos, glgf, gd):
    #this function convers SDR position into motif + offset format
    #return a dictionary of contacting residue pairs
    
    #hmmout = hmmer_call(pdbid, hmmdb, seq, eval)
    #outdict = {}
    split_line = 14 #why choose this ?
    for record in hmmout:
        for hit in record:
            for domain in hit:
                for key in dict_contact:
                    (tg_id, sc_id) = key
                    #tg_id = target residue number
                    #sc_id = source residue number
                    distance = dict_contact[key]
                    #print peppos, tg_id, tg_id == peppos
                    if tg_id == peppos:
                        #print tg_id, sc_id, domain.ali_from
                        if sc_id <= domain.ali_to and sc_id >= domain.ali_from:
                            motif_id, residue = get_offset(domain.target_sequence,
                                                  domain.hmm_sequence,
                                                  sc_id,
                                                  domain.ali_from,
                                                  domain.hmm_from,
                                                  glgf, gd)
                            if motif_id:
                                if outdict.has_key(motif_id):
                                 
                                    outdict[motif_id].append([hit.name,
                                                              residue,
                                                              distance])
                                else:
                                   
                                    outdict[motif_id] = [[hit.name,
                                                          residue,
                                                          distance]]
                        #else:
                        #    print "Error: This SDR is not in the alignment"
                    else:
                        pass
    #pprint.pprint(outdict)
    #print '--------------'
    return outdict
                    
def hmm_alignseq(_hmmout):
    #print 'hmm_aligseq start'
    _alignseq = ''
    
    for _record in _hmmout:
        for _hit in _record:
            for _domain in _hit:
                
                _pdbid = _hit.name
                _gaps = _domain.hmm_from - 1
                for i in range(_gaps):
                    _alignseq = _alignseq + '.'
                for (i, r) in enumerate(_domain.hmm_sequence):
                    if r != '.' and r != '-':
                        _alignseq = _alignseq + _domain.target_sequence[i]
    return _alignseq

def match_len(_maln):
    #Seq in alignment must be same length
    #adds '.' in the shorter seqs
    _max_length = 0
    for _seq in _maln:
        _max_length = max(_max_length, len(_seq.seq))
    for _seq in _maln:
        while len(_seq.seq) < _max_length:
            _seq.seq.append('.')
    return _maln
#print neighbors
#residue_list = Selection.unfold_entities(neighbors, 'R')
#print residue_list
###########################################################
def import_database(dbname, contactdat, gd, glgf):
    mydb = sqlite3.connect('database/'+dbname)
    curs = mydb.cursor()
    table_name = 'structure_contacts'
    sql = """
          create table if not exists motifoffset
          (motif varchar(4),
          offset interger);"""
    curs.execute(sql)
    curs.execute("insert into motifoffset values(?,?);", ("GD", gd))
    curs.execute("insert into motifoffset values(?,?);", ("GLGF", glgf))
    
    
    sql = """
          create table if not exists offset
          (peppos interger,
          motif varchar(4),
          offset interger,
          PDBID varchar(4),
          RES varchar(1),
          dist real);"""
    curs.execute(sql)
    
    sql = "insert into test_contact values(?,?,?,?,?,?);"
    for peppos in contactdat:
        for motif, offset in contactdat[peppos]:
            for pdbid, res, dist in contactdat[peppos][(motif,offset)]:
                dist = float(dist)
                curs.execute(sql, (peppos, motif, offset, pdbid, res, dist))
    
    mydb.commit()
    mydb.close()
    

#-------------
#Motif positions need to be recalculated if using different
#hmm profile


if __name__ == '__main__':
    glgf = 22
    gd = 55
    #-------------
    hmmdb = 'strucaln.hmm'
    
    #ccfile is a file of contacting info on which two chains in the
    #protein are interacting. (From Structure)
    
    ccfile= 'contactchain.txt'
    #ccfile = 'test.txt'
    entries = load_cc(ccfile)
    dict1 = {}
    final_out = {}
    multialn = []
    
    #Parser the ccfile
    for entry in entries:
        pdbid = entry[0]
        print pdbid
        source_chain = entry[1]
        target_chain = entry[2]
        parser = PDBParser()
        ppb = PPBuilder()
        try:
            structure = parser.get_structure('X', 'structure_data/'+pdbid+'.pdb')
        except(IOError):
            srcpath = PDBList().retrieve_pdb_file(pdbid)
            cmd = 'move'+" "+srcpath+" structure_data/"+pdbid+".pdb"
            #print cmd
            os.system(cmd)
            structure = parser.get_structure('X', 'structure_data/'+pdbid+'.pdb')
        
        
        #-----------------------------------------------
        #-----------------------------------------------
        if len(entry) == 3:
            #using the whole chain
            print "Entry has 3 columns, parsing all residues"
            structure = renum(structure, source_chain)
            structure = renum(structure, target_chain)
            dcontact = contact(structure, source_chain, target_chain)
            #print d
        elif len(entry) == 5:
            #need to crop source chain
            print "Entry has 5 columns, parsing:", entry[3], entry[4]
            structure = renum(structure, source_chain,
                              (int(entry[3]), int(entry[4])))
            structure = renum(structure, target_chain)
            dcontact = contact(structure, source_chain, target_chain)
            #print d
        else:
            print "Error entry"
        seqin = ''
        for pp in ppb.build_peptides(structure[0][source_chain]):
            seqin += pp.get_sequence()
            #print pdbid,seqin
        #print 'seqin:', dir(seqin)
        #hmmout = hmmer_call(pdbid, hmmdb, seqin, 0.01)
        for peppos in range(0,5):
            hmmout = hmmer_call(pdbid, hmmdb, seqin, 0.01)
            if final_out.has_key(peppos):
                dict1 = find_motif(final_out[peppos], hmmout,
                                   dcontact, peppos, glgf, gd)
            else:
                final_out[peppos] = {}
                dict1 = find_motif(final_out[peppos], hmmout,
                                   dcontact, peppos, glgf, gd)
            final_out[peppos] = dict1
            
        hmmout = hmmer_call(pdbid, hmmdb, seqin, 0.1)
        multialn.append(SeqRecord(Seq(hmm_alignseq(hmmout)), id = pdbid))
    multialn = match_len(multialn)
    multialn = MultipleSeqAlignment(multialn)
    #pprint(final_out)
    pprint.pprint(final_out)
    #write to output file 
    AlignIO.write(multialn, "tmp_aln.fasta", "fasta")
    import_database('pdzvalid.db', final_out, gd, glgf)
    #outfile = open('contactinfo.dat', 'wb')
    #pickle.dump([final_out, glgf, gd], outfile)
    #outfile.close()
    
    #for each in final_out:
    #    print each